package waf

import (
	"sync"

	"gorm.io/gorm"
)

type RuleManager struct {
	rules map[RuleType][]Rule
	mu    sync.RWMutex
	db    *gorm.DB
}

func NewRuleManager(db *gorm.DB) *RuleManager {
	rm := &RuleManager{
		rules: make(map[RuleType][]Rule),
		db:    db,
	}
	rm.LoadRules()
	return rm
}

// LoadRules 从数据库加载规则
func (rm *RuleManager) LoadRules() error {
	rm.mu.Lock()
	defer rm.mu.Unlock()

	var rules []Rule
	if err := rm.db.Where("status = ?", 1).Find(&rules).Error; err != nil {
		return err
	}

	newRules := make(map[RuleType][]Rule)
	for _, rule := range rules {
		newRules[rule.Type] = append(newRules[rule.Type], rule)
	}
	rm.rules = newRules
	return nil
}

// AddRule 添加规则
func (rm *RuleManager) AddRule(rule Rule) error {
	rm.mu.Lock()
	defer rm.mu.Unlock()

	if err := rm.db.Create(&rule).Error; err != nil {
		return err
	}

	if rule.Status == 1 {
		rm.rules[rule.Type] = append(rm.rules[rule.Type], rule)
	}
	return nil
}

// GetRules 获取指定类型的规则
func (rm *RuleManager) GetRules(ruleType RuleType) []Rule {
	rm.mu.RLock()
	defer rm.mu.RUnlock()
	return rm.rules[ruleType]
}
